#pragma once

#include <iosfwd>

#include <boost/cstdint.hpp>
#include <boost/scoped_array.hpp>
#include <boost/dynamic_bitset.hpp>
#include <boost/asio/ip/tcp.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/enable_shared_from_this.hpp>
#include <vector>

#include "torrent/fwd_decl.hpp"

namespace peer_wire {

enum MessageID
{
	CHOKE = 0, UNCHOKE, INTERESTED, NOT_INTERESTED, HAVE, BITFIELD, REQUEST, PIECE, CANCEL, PORT
	, EXTENDED = 20
	, NOT_SUPPORTED_MESSAGE_ID
};

struct Message
{
	boost::uint32_t _length;
	boost::uint8_t _msgID;

	Message();
	Message(boost::uint8_t msgID, boost::uint32_t length);
	virtual ~Message() {}
	std::size_t write(boost::asio::streambuf& buffer);
	std::size_t write(boost::asio::ip::tcp::socket& socket, boost::system::error_code& ec);
	boost::uint32_t getPayloadLength() const { return _length - sizeof(_msgID); }
	virtual void writePayload(std::ostream& os) = 0;
	virtual void readPayload(std::istream& is) = 0;
	virtual void dump(std::ostream& out);
	virtual void dispatch(torrent::P2PSessionPtr dispatcher) = 0;
	virtual void dispatchBySent(torrent::P2PSessionPtr dispatcher) = 0;
};

MessagePtr create_msg_from_socket(boost::asio::ip::tcp::socket& socket, boost::system::error_code& ec);
MessagePtr create_msg(boost::uint8_t msgID, boost::uint32_t msgLen);
std::string msg_str(boost::uint8_t msgID);

struct HandShake	: public boost::enable_shared_from_this<HandShake>
{
	std::string _pstr;
	std::string _infoHash;
	boost::array<char, 8> _reserved;
	std::string _peerID;

	HandShake();
	std::size_t write(boost::asio::streambuf& buffer);
	std::size_t write(boost::asio::ip::tcp::socket& socket, boost::system::error_code& ec);
	void read(boost::asio::streambuf& buffer, boost::uint8_t pstrlen);
	std::size_t read(boost::asio::ip::tcp::socket& socket, boost::system::error_code& ec);
	void dump(std::ostream& out);
	static boost::uint32_t getPayloadLengthFromPstrlen(boost::uint8_t pstrlen) { return pstrlen + 8 + 20 + 20; }
	void dispatch(torrent::P2PSessionPtr dispatcher);
	void dispatchBySent(torrent::P2PSessionPtr dispatcher);
};

struct KeepAlive	: public Message, public boost::enable_shared_from_this<KeepAlive>
{
	KeepAlive() : Message(NOT_SUPPORTED_MESSAGE_ID, 0) {}
	void writePayload(std::ostream& /*os*/) {}
	void readPayload(std::istream& /*is*/) {}
	void dump(std::ostream& out);
	void dispatch(torrent::P2PSessionPtr dispatcher);
	void dispatchBySent(torrent::P2PSessionPtr dispatcher);
};

struct Choke	: public Message, public boost::enable_shared_from_this<Choke>
{
	Choke() : Message(CHOKE, 1) {}
	void writePayload(std::ostream& /*os*/) {}
	void readPayload(std::istream& /*is*/) {}
	void dispatch(torrent::P2PSessionPtr dispatcher);
	void dispatchBySent(torrent::P2PSessionPtr dispatcher);
};

struct Unchoke	: public Message, public boost::enable_shared_from_this<Unchoke>
{
	Unchoke() : Message(UNCHOKE, 1) {}
	void writePayload(std::ostream& /*os*/) {}
	void readPayload(std::istream& /*is*/) {}
	void dispatch(torrent::P2PSessionPtr dispatcher);
	void dispatchBySent(torrent::P2PSessionPtr dispatcher);
};

struct Interested	: public Message, public boost::enable_shared_from_this<Interested>
{
	Interested() : Message(INTERESTED, 1) {}
	void writePayload(std::ostream& /*os*/) {}
	void readPayload(std::istream& /*is*/) {}
	void dispatch(torrent::P2PSessionPtr dispatcher);
	void dispatchBySent(torrent::P2PSessionPtr dispatcher);
};

struct NotInterested	: public Message, public boost::enable_shared_from_this<NotInterested>
{
	NotInterested() : Message(NOT_INTERESTED, 1) {}
	void writePayload(std::ostream& /*os*/) {}
	void readPayload(std::istream& /*is*/) {}
	void dispatch(torrent::P2PSessionPtr dispatcher);
	void dispatchBySent(torrent::P2PSessionPtr dispatcher);
};

struct Have	: public Message, public boost::enable_shared_from_this<Have>
{
	boost::uint32_t _pieceIndex;

	Have() : Message(HAVE, 5), _pieceIndex(0) {}
	Have(size_t pieceIdx) : Message(HAVE, 5), _pieceIndex(pieceIdx) {}
	void writePayload(std::ostream& os);
	void readPayload(std::istream& is);
	void dump(std::ostream& out);
	void dispatch(torrent::P2PSessionPtr dispatcher);
	void dispatchBySent(torrent::P2PSessionPtr dispatcher);
};

struct Bitfield	: public Message, public boost::enable_shared_from_this<Bitfield>
{
	boost::dynamic_bitset<> _bitfield;

	Bitfield(const boost::dynamic_bitset<>& bitfield);
	Bitfield(boost::uint32_t length);
	void writePayload(std::ostream& os);
	void readPayload(std::istream& is);
	void dump(std::ostream& out);
	void dispatch(torrent::P2PSessionPtr dispatcher);
	void dispatchBySent(torrent::P2PSessionPtr dispatcher);
};

struct Request	: public Message, public boost::enable_shared_from_this<Request>
{
	boost::uint32_t _index;
	boost::uint32_t _begin;
	boost::uint32_t _blockLength;

	Request() : Message(REQUEST, 13), _index(0), _begin(0), _blockLength(0) {}
	void writePayload(std::ostream& os);
	void readPayload(std::istream& is);
	void dump(std::ostream& out);
	void dispatch(torrent::P2PSessionPtr dispatcher);
	void dispatchBySent(torrent::P2PSessionPtr dispatcher);
};

struct Piece	: public Message, public boost::enable_shared_from_this<Piece>
{
	boost::uint32_t _index;
	boost::uint32_t _begin;
	boost::scoped_array<boost::uint8_t> _block;

	Piece(boost::uint32_t length);
	Piece(boost::uint32_t index, boost::uint32_t begin, boost::uint32_t blockLength);
	boost::uint32_t getBlockLength() { return getPayloadLength() - sizeof(_index) - sizeof(_begin); }
	void setBlockLength(boost::uint32_t blockLength) { _length = sizeof(_index) + sizeof(_begin) + blockLength; }
	void writePayload(std::ostream& os);
	void readPayload(std::istream& is);
	void dump(std::ostream& out);
	void dispatch(torrent::P2PSessionPtr dispatcher);
	void dispatchBySent(torrent::P2PSessionPtr dispatcher);
};

struct Cancel	: public Message, public boost::enable_shared_from_this<Cancel>
{
	boost::uint32_t _index;
	boost::uint32_t _begin;
	boost::uint32_t _blockLength;

	Cancel() : Message(CANCEL, 13), _index(0), _begin(0), _blockLength(0) {}
	void writePayload(std::ostream& os);
	void readPayload(std::istream& is);
	void dump(std::ostream& out);
	void dispatch(torrent::P2PSessionPtr dispatcher);
	void dispatchBySent(torrent::P2PSessionPtr dispatcher);
};

struct Port	: public Message, public boost::enable_shared_from_this<Port>
{
	boost::uint16_t _port;

	Port() : Message(PORT, 3), _port(0) {}
	void writePayload(std::ostream& os);
	void readPayload(std::istream& is);
	void dump(std::ostream& out);
	void dispatch(torrent::P2PSessionPtr dispatcher);
	void dispatchBySent(torrent::P2PSessionPtr dispatcher);
};

// TODO : extension. http://bittorrent.org/beps/bep_0010.html
// TODO : extension. http://www.rasterbar.com/products/libtorrent/extension_protocol.html
struct Extended	: public Message, public boost::enable_shared_from_this<Extended>
{
	boost::uint8_t _extendedMsgID;
	be_node* _extendedDict;

	Extended(boost::uint32_t length);
	~Extended();
	boost::uint32_t getDictLength() { return getPayloadLength() - sizeof(_extendedMsgID); }
	void writePayload(std::ostream& os);
	void readPayload(std::istream& is);
	void dump(std::ostream& out);
	void dispatch(torrent::P2PSessionPtr dispatcher);
	void dispatchBySent(torrent::P2PSessionPtr dispatcher);
};

} // namespace peer_wire
